import cv2
import time
import os
from yoloSegV11.rknnpool import yoloPoolExecutor
# 图像处理函数，实际应用过程中需要自行修改
# from yolo import yolo_detect
import numpy as np

# 创建输出目录
output_dir = 'output'
os.makedirs(output_dir, exist_ok=True)
# 读取视频
cap = cv2.VideoCapture('./data/1.avi')

######### yolo #########
poseModelPath = "./models/merge_pose_epoch_600-rk3588-fp16.rknn"
# 线程数, 增大可提高帧率
poseNpuIDs = [0]
# 初始化rknn池
posePool = yoloPoolExecutor(
    rknnModel=poseModelPath,
    npu_ids=poseNpuIDs,
    func=yolo_detect)

poseModelPath = "./models/merge_seg_epoch201-rk3588-fp16.rknn"
# 线程数, 增大可提高帧率
segNpuIDs = [1]
# 初始化rknn池
segPool = yoloPoolExecutor(
    rknnModel=segModelPath,
    npu_ids=segNpuIDs,
    func=yolo_detect)

######### samEncoder #########
samEncoderModelPath = "./rknnModel/mobilesam_encoder.rknn"
samEncoderNpuID = [1]
samEncoderPool = samEncoderPoolExecutor(
    rknnModel=samEncoderModelPath,
    npu_ids=samEncoderNpuID,
    func=sam_encoder)

######### samDecoder #########
samDecoderModelPath = "./rknnModel/mobilesam_decoder.rknn"
samDecoderNpuID = [2]
samDecoderPool = samDecoderPoolExecutor(
    rknnModel=samDecoderModelPath,
    npu_ids=samDecoderNpuID,
    func=sam_decoder)
# 初始化异步所需要的帧
if (cap.isOpened()):
    # 跳过前500帧
    for i in range(500):
        ret, frame = cap.read()
        if not ret:
            cap.release()
            del yoloPool
            exit(-1)
    
    # 初始化异步处理所需的帧
    for i in range(1):
        ret, frame = cap.read()
        if not ret:
            cap.release()
            del yoloPool
            exit(-1)
        yoloPool.put(frame)

frames, loopTime, initTime = 0, time.time(), time.time()

# 添加计时器变量
yolo_time = 0
sam_encoder_time = 0
sam_decoder_time = 0
total_frames = 0

computeAngle = ComputeAngle()

def draw(images, masks, init_coords, center, angle, quad_points, color=(144, 144, 30)):
    """
    在图像上绘制掩码、中心点、角度和四边形顶点
    
    参数:
        images: 原始图像
        masks: 分割掩码
        init_coords: 初始边界框坐标
        center: 中心点坐标
        angle: 角度值
        quad_points: 四边形顶点坐标
        color: 掩码颜色，默认为(144, 144, 30)
    """
    print(f"初始坐标: {init_coords}")
    
    # 确保掩码是二维的
    if len(masks.shape) > 2:
        masks = masks.squeeze()
    
    # 检查并调整掩码尺寸以匹配图像
    h, w = images.shape[:2]
    if masks.shape != (h, w):
        print(f"调整掩码尺寸从 {masks.shape} 到 {(h, w)}")
        masks = cv2.resize(masks.astype(np.uint8), (w, h), interpolation=cv2.INTER_NEAREST)
    
    alpha = 0.5
    color = np.array(color)
    
    # 创建彩色掩码图像
    mask_image = np.zeros((h, w, 3), dtype=np.uint8)
    mask_image[masks > 0] = color
    
    # 将掩码叠加到原始图像上
    overlay = cv2.addWeighted(images, 1, mask_image, alpha, 0)
    images = np.where(mask_image != 0, overlay, images)
    
    # 绘制中心点和角度
    if center is not None:
        center = tuple(map(int, center))
        cv2.circle(images, center, 5, (0, 0, 255), -1)
        cv2.putText(images, f"Angle: {angle:.2f}°", 
                   (center[0], center[1] - 10), 
                   cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2)
    
    # 绘制四边形顶点
    if quad_points is not None:
        for point in quad_points:
            point = tuple(map(int, point))
            cv2.circle(images, point, 5, (0, 255, 0), -1)
    
    # 保存结果图像
    cv2.imwrite("result_image.jpg", images)
    
    return images

while (cap.isOpened()):
    frames += 1
    total_frames += 1
    ret, frame = cap.read()
    if not ret:
        break
    
    # YOLO处理计时
    yolo_start = time.time()
    yoloPool.put(frame)
    result, flag = yoloPool.get()
    yolo_time += time.time() - yolo_start
    
    # 获取检测框
    box = result['max_box'][0]
    # print(box)
    
    # 裁剪图像
    cropped_frame, new_box, crop_box = computeAngle.crop_frame(frame, box)
    
    # 先检查是否可以复用上一帧的结果
    result_dict = computeAngle.process_detection(cropped_frame, new_box)
    
    # 只有在不能复用结果时才执行SAM处理
    if not result_dict['reused']:
        # SAM编码器处理计时
        sam_encoder_start = time.time()
        samEncoderPool.put(cropped_frame)
        img_embeds, flag = samEncoderPool.get()
        sam_encoder_time += time.time() - sam_encoder_start

        # SAM解码器处理计时
        sam_decoder_start = time.time()
        samDecoderPool.put(img_embeds, cropped_frame.shape[:2], new_box)
        sam_mask, flag = samDecoderPool.get()
        sam_decoder_time += time.time() - sam_decoder_start

        # 使用新的掩码更新结果
        result_dict = computeAngle.process_detection(cropped_frame, new_box, sam_mask)
    else:
        # 如果复用了结果，使用上一帧的掩码
        sam_mask = result_dict['mask']
        # 更新计时器，复用时不计算SAM时间
        sam_encoder_time += 0
        sam_decoder_time += 0
    
    if result_dict['success']:
        center = result_dict['center']
        angle = result_dict['angle']
        quad_points = result_dict['quad_points']
        
        # 在裁剪后的图像上绘制结果
        processed_frame = draw(cropped_frame, sam_mask, new_box, center, angle, quad_points)
        
        # 将处理后的结果合并回原始图像
        result_frame = computeAngle.merge_result(frame, processed_frame, crop_box, center, angle, quad_points)
        
        # 添加复用信息到图像
        if result_dict['reused']:
            cv2.putText(result_frame, f"Reused: {computeAngle.reuse_counter}/10", 
                       (10, 90), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 255), 2)
    else:
        result_frame = frame
        cv2.putText(result_frame, "处理失败", 
                   (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
    
    print(f"第{frames}帧完成")
    print(f"YOLO耗时: {yolo_time/total_frames:.3f}s, SAM编码器耗时: {sam_encoder_time/total_frames:.3f}s, SAM解码器耗时: {sam_decoder_time/total_frames:.3f}s")
    
    # 保存每一帧到output目录
    out_path = os.path.join(output_dir, f'frame_{frames:06d}.jpg')
    cv2.imwrite(out_path, result_frame)

    if frames % 30 == 0:
        print("30帧平均帧率:\t", 30 / (time.time() - loopTime), "帧")
        print(f"平均处理时间 - YOLO: {yolo_time/total_frames:.3f}s, SAM编码器: {sam_encoder_time/total_frames:.3f}s, SAM解码器: {sam_decoder_time/total_frames:.3f}s")
        loopTime = time.time()

print("总平均帧率\t", frames / (time.time() - initTime))
print(f"最终平均处理时间 - YOLO: {yolo_time/total_frames:.3f}s, SAM编码器: {sam_encoder_time/total_frames:.3f}s, SAM解码器: {sam_decoder_time/total_frames:.3f}s")
# 释放cap和rknn线程池
cap.release()
# cv2.destroyAllWindows()
yoloPool.release()
